/*
 *   Copyright 2023 <a href="mailto:asialjim@hotmail.com">Asial Jim</a>
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */
package com.asialjim.encommunicate.rsa.filter;

import com.asialjim.encommunicate.base.rsa.key.RSAServerKey;
import com.asialjim.encommunicate.base.rsa.util.RSAServerUtil;
import com.asialjim.encommunicate.rsa.management.key.RSAKeyLookUpComponent;
import com.asialjim.encommunicate.rsa.management.pojo.RSAKeyResult;
import com.google.gson.Gson;
import jakarta.servlet.*;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.http.InvalidMediaTypeException;
import org.springframework.http.MediaType;
import org.springframework.http.converter.json.GsonBuilderUtils;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Objects;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * 基于 RSA算法的过滤器
 *
 * @author Copyright © <a href="mailto:asialjim@hotmail.com">Asial Jim</a>   Co., LTD
 * @version 1.0
 * @since 2023/2/27, &nbsp;&nbsp; <em>version:1.0</em>, &nbsp;&nbsp; <em>java version:17</em>
 */
@Slf4j
public class SimpleRSAFilter implements Filter {
    private static final Gson gson = GsonBuilderUtils.gsonBuilderWithBase64EncodedByteArrays().create();
    public static final String APPID_HEADER_NAME = "appid";
    public static final String SIGNATURE_HEADER_NAME = "signature";
    private static final String[] APPLICATION = {"application"};
    private static final String[] TEXT = {"text", "json", "xml"};
    private static final String[] MEDIA = {"media", "image", "multipart", "video", "audio"};
    private static final String[] STATIC_RESOURCE = {".html", ".xhtml", ".shtml", ".css", ".js"};
    private final RSAKeyLookUpComponent rsaKeyLookUpComponent;
    private final ApplicationEnDecryptChecker applicationEnDecryptChecker;
    public static final SimpleRSAFilter NON_RSA_FILTER = new SimpleRSAFilter(null){
        @Override
        public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
            chain.doFilter(request,response);
        }
    };

    public SimpleRSAFilter(RSAKeyLookUpComponent rsaKeyLookUpComponent) {
        this(rsaKeyLookUpComponent,null);
    }

    public SimpleRSAFilter(RSAKeyLookUpComponent rsaKeyLookUpComponent, ApplicationEnDecryptChecker applicationEnDecryptChecker) {
        this.rsaKeyLookUpComponent = rsaKeyLookUpComponent;
        if (Objects.isNull(applicationEnDecryptChecker)) 
            applicationEnDecryptChecker = ApplicationEnDecryptChecker.DEFAULT;
        
        this.applicationEnDecryptChecker = applicationEnDecryptChecker;
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        // 非 HTTP SERVLET 请求或响应，直接pass
        if (!(request instanceof HttpServletRequest httpServletRequest) || !(response instanceof HttpServletResponse httpServletResponse)) {
            chain.doFilter(request, response);
            return;
        }

        String appid = httpServletRequest.getHeader(APPID_HEADER_NAME);

        HttpServletRequest targetHttpServletRequest = httpServletRequest;
        SimpleRSAHttpServletResponseWrapper targetResponse = null;
        RSAServerKey rsaServerKey = null;
        RSAServerUtil rsaServerUtil = null;

        try {
            // 判断请求是否启用加密通讯
            boolean requestEncrypt = requestEncryptEnable(httpServletRequest, appid);
            
            // 判断响应是否启用加密通讯
            boolean responseEncrypt = responseEncryptEnable(httpServletRequest,appid);

            // 需要开启请求加密
            if (requestEncrypt) {
                rsaServerKey = rsaKeyLookUpComponent.serverKeyThrowable(appid);
                rsaServerUtil = new RSAServerUtil(rsaServerKey);
                targetHttpServletRequest = new SimpleRSAHttpServletRequestWrapper(httpServletRequest, rsaServerUtil);
            }

            // 需要开启响应加密
            if (responseEncrypt) {
                if (Objects.isNull(rsaServerKey)) rsaServerKey = rsaKeyLookUpComponent.serverKeyThrowable(appid);
                if (Objects.isNull(rsaServerUtil)) rsaServerUtil = new RSAServerUtil(rsaServerKey);
                
                // 包装响应，并开启处理业务逻辑
                targetResponse = new SimpleRSAHttpServletResponseWrapper(httpServletResponse);

                // 对响应结果进行加密
                wrapEncryptResponse(chain, httpServletResponse, targetHttpServletRequest, targetResponse, rsaServerUtil);
            }
            
            // 不需要对响应加密
            else chain.doFilter(targetHttpServletRequest, httpServletResponse);

        } catch (InvalidMediaTypeException invalidMediaTypeException) {
            
            // 不支持的 Content Type
            sendResponseException(415, "Unsupported Media Type", invalidMediaTypeException, httpServletResponse); // 请求内容不支持
        } catch (IllegalStateException illegalStateException) {
            
            // 加解密错误
            sendResponseException(403, "Request Forbidden", illegalStateException, httpServletResponse);
        } finally {
            if (Objects.nonNull(targetResponse))
                targetResponse.close();
        }
    }

    private static void wrapEncryptResponse(FilterChain chain, HttpServletResponse httpServletResponse, HttpServletRequest targetHttpServletRequest,
                                            SimpleRSAHttpServletResponseWrapper targetResponse, RSAServerUtil rsaServerUtil) throws IOException, ServletException {
        chain.doFilter(targetHttpServletRequest, targetResponse);

        // 获取业务响应内容
        byte[] content = targetResponse.getContent();

        // 对响应内容进行加密
        byte[] encrypt = rsaServerUtil.encrypt(content);

        // 对响应结果进行加签
        String sign = rsaServerUtil.sign(content);

        // 拦截重写加密后的数据
        httpServletResponse.setHeader(SIGNATURE_HEADER_NAME, sign);

        // 重写内容长度
        httpServletResponse.setContentLength(encrypt.length);

        // 重写加密后的响应内容
        ServletOutputStream outputStream = httpServletResponse.getOutputStream();
        outputStream.write(encrypt);
        outputStream.flush();
    }

    /**
     * 向 response 写异常
     *
     * @param status              {@link Integer 网络状态码}
     * @param msg                 {@link String 响应信息}
     * @param e                   {@link Throwable 异常}
     * @param httpServletResponse {@link HttpServletResponse 响应}
     * @since 2023/2/28
     */
    private static void sendResponseException(int status, String msg, Throwable e, HttpServletResponse httpServletResponse) throws IOException {
        RSAKeyResult<Object> res = RSAKeyResult.builder().code(String.valueOf(status)).msg(msg).errList(Stream.of(e.getMessage()).collect(Collectors.toList())).build();
        writeResponse(httpServletResponse, status, gson.toJson(res));
    }

    /**
     * 判断 appid 请求是否启用加密通讯
     *
     * @param appid {@link String appid}
     * @return {@link Boolean }
     * @since 2023/2/28
     */
    private boolean appEncNeed(String appid) {
        return applicationEnDecryptChecker.reqEncNeed(appid);
    }

    /**
     * 判断响应结果是否需要加密
     *
     * @param appid {@link String appid}
     * @return {@link Boolean 是否需要加密}
     * @since 2023/2/27
     */
    private boolean responseEncryptEnable(HttpServletRequest request,String appid) {
        String requestURI = request.getRequestURI();

        // 静态资源
        if (StringUtils.containsAny(requestURI, STATIC_RESOURCE)) return false;

        return applicationEnDecryptChecker.resEncNeed(appid);
    }

    /**
     * 判断是否需要解密
     * <br/> 当不需要解密时，返回 true
     *
     * @param request {@link HttpServletRequest request}
     * @return {@link Boolean 不需要解密时，返回True， 否则返回 False}
     * @since 2023/2/27
     */
    private boolean requestEncryptEnable(HttpServletRequest request, String appid) {
        String requestURI = request.getRequestURI();

        // 静态资源
        if (StringUtils.containsAny(requestURI, STATIC_RESOURCE)) return false;

        // 应用未启用请求加密通讯
        if (!appEncNeed(appid)) return false;

        // 应用启用加密通讯的情况下
        // 请求无请求体，无需加密通讯
        if (request.getContentLength() <= 0 || request.getContentLengthLong() <= 0) return false;

        // 获取请求体媒体类型
        MediaType mediaType = MediaType.valueOf(request.getContentType());
        String type = mediaType.getType();
        String subtype = mediaType.getSubtype();

        // 请求体是图片，声音，视频等媒体文件，不需要加密通讯
        if (StringUtils.equalsAnyIgnoreCase(type, MEDIA)) return false;

        // 请求体类型是文本类型，需要加密通讯
        if (StringUtils.equalsAnyIgnoreCase(type, TEXT)) return true;

        // 请求体类型是应用类型
        if (StringUtils.equalsAnyIgnoreCase(type, APPLICATION)) {
            // 子类型为文本类型时，需要开启加密通讯
            return StringUtils.containsAnyIgnoreCase(subtype, TEXT);
        }

        // 其他类型，不支持
        throw new InvalidMediaTypeException(request.getContentType(), "不支持的 Content Type");
    }

    /**
     * 向 response 中写 json 响应数据
     *
     * @param response   {@link HttpServletResponse response}
     * @param status     {@link Integer status}
     * @param resultJson {@link String resultJson}
     * @since 2023/2/28
     */
    private static void writeResponse(HttpServletResponse response, int status, String resultJson) throws IOException {
        ServletOutputStream outputStream = null;
        try {
            response.setStatus(status);
            response.setHeader("Content-Type", "application/json;charset=utf-8");

            byte[] illegalRequestBytes = resultJson.getBytes(StandardCharsets.UTF_8);
            response.setHeader("Content-Length", String.valueOf(illegalRequestBytes.length));
            outputStream = response.getOutputStream();
            outputStream.write(illegalRequestBytes);
            outputStream.flush();
        } finally {
            IOUtils.closeQuietly(outputStream);
        }
    }
}